/*
 * License: Copyright (c) Huawei Technologies Co., Ltd. 2012-2019. All rights reserved.
 * Description: main function, launch the app.
 * Date: 2019-11-30 10:57:41
 * LastEditTime : 2020-02-17 12:19:14
 */

#include <unistd.h>
#include <thread>
#include <condition_variable>
#include <sstream>
#include <string>
#include "hiaiengine/api.h"
#include "proto/graph_config.pb.h"
#include "config_parser.h"
#include "str_num.h"
#include "trans_data.h"
#include "recv_data.h"

using namespace std;

const int BASE_LINE = 100;
const int SLEEP_TIME = 500;

const int MIN_BIT_RATE = 10;
const int MAX_BIT_RATE = 28672;
const int MIN_FRAME_RATE = 1;
const int MAX_FRAME_RATE = 240;

bool g_isFinished = false;
bool g_isReady = false;
static mutex g_mtx;
static condition_variable g_cv;

static const uint32_t MAX_SLEEP_TIMER = 24 * 60 * 60;  // 24h

uint64_t g_channelNum = 0;

vector<string> g_videos;
vector<string> g_pushStreams;

const string g_srcEngineName = "SrcEngine";
const string g_dstEngineName = "DestEngine";

const int VBR = 1;
const int CBR = 2;

struct ConfigInfo {
    int graphId;
    int srcEngineId;
    int destEngineId;
    string graphPath;
    int chipId;
    int channelNum;
    int outputFormat;
    int vencFrameRate;
    int vencBitRate;
    int vencCtlType;

    ConfigInfo()
    {
        graphId = 0;
        srcEngineId = 0;
        destEngineId = 0;
        graphPath = "";
        chipId = 0;
        channelNum = 0;
        outputFormat = 0;
        vencFrameRate = 0;
        vencBitRate = 0;
        vencCtlType = 0;
    };
};

void CheckAllProcessFinished()
{
    while (!g_isFinished) {
        this_thread::sleep_for(chrono::milliseconds(15));
    }

    // all channel are finished, notify main thread to destroy the graphs.
    unique_lock <mutex> lck(g_mtx);
    g_isReady = true;
    g_cv.notify_all();
}

// init the graphs on mutil chips
HIAI_StatusT InitAndStartGraph(string graphPath, int graphId, int destEngineId, int channelNum, int chipId)
{
    HIAI_StatusT status;
    hiai::GraphConfigList graphConfigList;
    // parse the graph file
    status = hiai::Graph::ParseConfigFile(graphPath, graphConfigList);
    if (status != HIAI_OK) {
        HIAI_ENGINE_LOG("Failed to parse graph file!");
        cout << "Failed to parse graph file!" << endl;
        cout << graphPath << endl;
        return status;
    }

    for (int i = 0; i < channelNum; ++i) {
        int deviceId = chipId;
        int curGraphId = graphId + chipId * BASE_LINE + i;

        // init the hiaiengine
        status = HIAI_Init(deviceId);
        if (status != HIAI_OK) {
            HIAI_ENGINE_LOG("HIAI_Init failed!");
            cout << "HIAI_Init failed! Please check the npu device or driver!" << endl;
            return status;
        }

        // modify the graph file
        hiai::GraphConfig* graphCfg = graphConfigList.mutable_graphs(0);
        graphCfg->set_device_id(to_string(deviceId));
        graphCfg->set_graph_id(curGraphId);

        // create the graph according to the prototxt
        status = hiai::Graph::CreateGraph(graphConfigList);
        if (status != HIAI_OK) {
            HIAI_ENGINE_LOG("Failed to create the graph!");
            cout << "Failed to create the graph!" << endl;
            return status;
        }

        pid_t pid = getpid();
        HIAI_ENGINE_LOG(HIAI_IDE_ERROR, "PID:%d GraphID:%d", pid, curGraphId);

        // get the graph instance
        shared_ptr<hiai::Graph> graph = hiai::Graph::GetInstance(curGraphId);
        if (graph == nullptr) {
            HIAI_ENGINE_LOG("Failed to get the graph instance!");
            cout << "Failed to get the graph instance!" << endl;
            return status;
        }

        // config the destination engine and the port
        hiai::EnginePortID targetEngine;
        targetEngine.graph_id = curGraphId;
        targetEngine.engine_id = destEngineId; // destEngine
        targetEngine.port_id = 0;

        graph->SetDataRecvFunctor(targetEngine,
            shared_ptr<UserDataRecvInterface>(new UserDataRecvInterface("start")));

        HIAI_ENGINE_LOG(HIAI_IDE_ERROR, "DEVICE:%d has been started", deviceId);
    }

    return HIAI_OK;
}

int ParamCheck(ConfigInfo& cfgInfo)
{
    if (cfgInfo.channelNum < 0) {
        cout << "channel num is incorrect, please check it!" << endl;
        return -1;
    }

    if (cfgInfo.chipId < 0) {
        cout << "chip id is incorrect, please check it!" << endl;
        return -1;
    }

    if (cfgInfo.outputFormat != 0 && cfgInfo.outputFormat != 1) {
        cout << "output Format is incorrect, please check it! Only 0 or 1." << endl;
        return -1;
    }

    if (cfgInfo.vencBitRate < MIN_BIT_RATE || cfgInfo.vencBitRate > MAX_BIT_RATE) {
        cout << "vencBitRate([10, 28672]) is incorrect, please check it!" << endl;
        return -1;
    }

    if (cfgInfo.vencFrameRate < MIN_FRAME_RATE || cfgInfo.vencFrameRate > MAX_FRAME_RATE) {
        cout << "vencFrameRate([1, 240]) is incorrect, please check it!" << endl;
        return -1;
    }

    if (cfgInfo.vencCtlType != VBR && cfgInfo.vencCtlType != CBR) {
        cout << "vencCtlType is incorrect, please check it! Only 1 or 2." << endl;
        return -1;
    }

    return 0;
}

int CheckConfigInfoByProto(string graphPath, ConfigInfo& cfgInfo)
{
    HIAI_StatusT status;
    hiai::GraphConfigList graphConfigList;
    // parse the graph file
    status = hiai::Graph::ParseConfigFile(graphPath, graphConfigList);
    if (status != HIAI_OK) {
        HIAI_ENGINE_LOG("Failed to parse graph file!");
        cout << "Failed to parse graph file!" << endl;
        cout << graphPath << endl;
        return -1;
    }

    hiai::GraphConfig graphCfg = graphConfigList.graphs(0);
    int graphIdFromProto = graphCfg.graph_id();
    if (graphIdFromProto != cfgInfo.graphId) {
        cout << "graph id is incorrect! please check it!" << endl;
        return -1;
    }

    int srcEngineId = -1;
    int dstEngineId = -1;
    int engineNum = graphCfg.engines_size();
    for (int i = 0; i < engineNum; ++i) {
        hiai::EngineConfig engineCfg = graphCfg.engines(i);
        if (engineCfg.engine_name() == g_srcEngineName) {
            srcEngineId = engineCfg.id();
        }
        if (engineCfg.engine_name() == g_dstEngineName) {
            dstEngineId = engineCfg.id();
        }
    }

    if (srcEngineId != cfgInfo.srcEngineId) {
        cout << "src engine id is incorrect! please check it!" << endl;
        return -1;
    }

    if (dstEngineId != cfgInfo.destEngineId) {
        cout << "dest engine id is incorrect! please check it!" << endl;
        return -1;
    }

    return 0;
}

int main(int argc, char* argv[])
{
    if (argc != 2) {
        cout << "Command format is incorrect!" << endl \
                  << "./MediaCodec config.ini" << endl;
        return -1;
    }

    // parse ini file
    map<string, string> iniCfg;
    ReadConfig(argv[1], iniCfg);
    PrintConfig(iniCfg);

    int graphId;
    int srcEngineId;
    int destEngineId;
    string graphPath;
    int chipId;
    int channelNum;
    int outputFormat;
    int vencFrameRate;
    int vencBitRate;
    int vencCtlType;

    try {
        graphId = str2num(iniCfg.at("graph_id"));
        srcEngineId = str2num(iniCfg.at("src_engine_id"));
        destEngineId = str2num(iniCfg.at("dest_engine_id"));
        graphPath = iniCfg.at("graph_path");
        chipId = str2num(iniCfg.at("chip_id"));
        channelNum = str2num(iniCfg.at("channel_num"));
        outputFormat = str2num(iniCfg.at("output_video_format"));
        vencFrameRate = str2num(iniCfg.at("venc_frame_rate"));
        vencBitRate = str2num(iniCfg.at("venc_bit_rate"));
        vencCtlType = str2num(iniCfg.at("venc_control_bitrate_type"));
    } catch (std::out_of_range& e) {
        cout << " [ERROR]" << "Some key-value is incorrect, please check the config file!" << endl;
        return -1;
    }


    g_channelNum = channelNum;

    ConfigInfo cfgInfo;
    cfgInfo.graphId = graphId;
    cfgInfo.srcEngineId = srcEngineId;
    cfgInfo.destEngineId = destEngineId;
    cfgInfo.graphPath = graphPath;
    cfgInfo.chipId = chipId;
    cfgInfo.channelNum = channelNum;
    cfgInfo.outputFormat = outputFormat;
    cfgInfo.vencBitRate = vencBitRate;
    cfgInfo.vencFrameRate = vencFrameRate;
    cfgInfo.vencCtlType = vencCtlType;

    int res = ParamCheck(cfgInfo);
    if (res == -1) {
        return -1;
    }

    res = CheckConfigInfoByProto(graphPath, cfgInfo);
    if (res == -1) {
        return -1;
    }

    for (int i = 0; i < channelNum; ++i) {
        string prefix = "video";
        string channelName = prefix + to_string(i);
        if (iniCfg[channelName] == "") {
            continue;
        }
        g_videos.push_back(iniCfg[channelName]);
    }

    for (int i = 0; i < channelNum; ++i) {
        string prefix = "stream";
        string streamName = prefix + to_string(i);
        if (iniCfg[streamName] == "") {
            continue;
        }
        g_pushStreams.push_back(iniCfg[streamName]);
    }

    if (channelNum != g_videos.size() || channelNum != g_pushStreams.size()) {
        cout << "the video num and stream num must equal to the channel num!" << endl;
        return -1;
    }

    // init and create the graphs on mutil chips
    HIAI_StatusT status = HIAI_OK;
    status = InitAndStartGraph(graphPath, graphId, destEngineId, channelNum, chipId);
    if (status != HIAI_OK) {
        HIAI_ENGINE_LOG("Failed to start graphs!");
        cout << "Failed to start graphs!" << endl;
        return -1;
    }

    // sent data
    for (int i = 0; i < channelNum; ++i) {
        int deviceId = chipId;
        int curGraphId = graphId + chipId * BASE_LINE + i;

        shared_ptr<hiai::Graph> graph = hiai::Graph::GetInstance(curGraphId);
        if (graph == nullptr) {
            HIAI_ENGINE_LOG("Failed to get the graph instance!");
            cout << "Failed to get the graph instance!" << endl;
            return -1;
        }

        hiai::EnginePortID srcEngine;
        srcEngine.graph_id = curGraphId;
        srcEngine.engine_id = srcEngineId; // srcEngine
        srcEngine.port_id = 0;

        shared_ptr<TransData> transData = shared_ptr<TransData>(new TransData);
        transData->graphId = curGraphId;
        transData->deviceId = deviceId;
        transData->channelId = i;

        // output format setting
        transData->outputFormat = outputFormat;
        if (outputFormat == 0) {
            cout << g_videos[i] << " OUTPUT_FORMAT: " << "D1" << endl;
        } else {
            cout << g_videos[i] << " OUTPUT_FORMAT: " << "CIF" << endl;
        }

        transData->vencFrameRate = vencFrameRate;
        transData->vencBitRate = vencBitRate;
        transData->vencCtlType = vencCtlType;
        graph->SendData(srcEngine, "TransData", static_pointer_cast<void>(transData));
        this_thread::sleep_for(chrono::milliseconds(SLEEP_TIME));
    }

    // wait the processes to be finished
    thread checkThread(CheckAllProcessFinished);
    checkThread.join();
    unique_lock <mutex> lck(g_mtx);
    g_cv.wait_for(lck, chrono::seconds(MAX_SLEEP_TIMER), [] { return g_isReady; });

    // destroy the graphs
    pid_t pid = getpid();
    for (int i = 0; i < channelNum; ++i) {
        int curGraphId = graphId + chipId * BASE_LINE + i;
        HIAI_ENGINE_LOG(HIAI_IDE_ERROR, "PID:%d, DestroyedGraphID: %d", pid, curGraphId);
        hiai::Graph::DestroyGraph(curGraphId);
    }

    return 0;
}